import uvicorn
from fastapi import FastAPI
from fastapi.responses import FileResponse
import json
import matplotlib.pyplot as plt
from docx import Document
from docx.shared import Inches
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.oxml.ns import qn
from matplotlib import rcParams
from docx.shared import Pt, Inches, RGBColor
import os
from typing import Any
from openai import OpenAI
import copy
from docx.enum.table import WD_TABLE_ALIGNMENT
from docx.oxml.shared import OxmlElement
from docx.oxml import parse_xml
import numpy as np
from datetime import datetime
from docx2pdf import convert
import tempfile

from pydantic import BaseModel
from starlette.middleware.cors import CORSMiddleware

# # 从本地config配置文件读取api_key，初始化openai client
# with open('apikey', 'r') as f:
#     api_key = f.read()
#     client = OpenAI(
#         # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
#         api_key=api_key,
#         base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
#     )
#     f.close()
api_key = "sk-f778cacdf67445c595442664cfa063fd"
client = OpenAI(
    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
    api_key=api_key,
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

def create_beautiful_charts(json_data):
    """创建美观的图表"""
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
    plt.rcParams['axes.unicode_minus'] = False
    
    # 1. 生成美观的 all_score 分布图
    plt.figure(figsize=(12, 6))
    colors = plt.cm.viridis(np.linspace(0, 1, len(json_data["score_all"])))
    bars = plt.bar(range(len(json_data["score_all"])), json_data["score_all"], 
                   color=colors, alpha=0.8, edgecolor='white', linewidth=0.5)
    
    # 添加数值标签
    for i, bar in enumerate(bars):
        height = bar.get_height()
        if height > 0:
            plt.text(bar.get_x() + bar.get_width()/2, height + 0.05, 
                    str(int(height)), ha='center', va='bottom', fontsize=8)
    
    plt.title("各题得分分布图", fontsize=16, fontweight='bold', pad=20)
    plt.xlabel("题目序号", fontsize=12)
    plt.ylabel("得分", fontsize=12)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig("all_score.png", dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()

    # 2. 生成美观的 dimensions 维度图
    dim_names = [d["name"] for d in json_data["dimensions"]]
    dim_scores = [d["score"] for d in json_data["dimensions"]]
    
    plt.figure(figsize=(10, 6))
    colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4', '#FFEAA7']
    bars = plt.bar(dim_names, dim_scores, color=colors[:len(dim_scores)], 
                   alpha=0.8, edgecolor='white', linewidth=1)
    
    # 添加数值标签和百分比
    total_score = sum(dim_scores)
    for bar, score in zip(bars, dim_scores):
        height = bar.get_height()
        percentage = (score / total_score) * 100 if total_score > 0 else 0
        plt.text(bar.get_x() + bar.get_width()/2, height + 0.1, 
                f'{score}\n({percentage:.1f}%)', 
                ha='center', va='bottom', fontsize=10, fontweight='bold')
    
    plt.title("各维度得分情况", fontsize=16, fontweight='bold', pad=20)
    plt.ylabel("得分", fontsize=12)
    # plt.xticks(rotation=45, ha='right')
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig("dimensions.png", dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()

def add_element_with_style(doc, element_type, text, font_size=12, font_name='宋体', 
                          bold=False, color=None, alignment=WD_PARAGRAPH_ALIGNMENT.LEFT):
    """添加带样式的元素"""
    if element_type == 'heading':
        heading = doc.add_heading("", level=1)
        run = heading.add_run(text)
        run.font.name = font_name
        run._element.rPr.rFonts.set(qn('w:eastAsia'), font_name)
        run.font.size = Pt(font_size)
        run.font.bold = bold
        if color:
            run.font.color.rgb = RGBColor(*color)
        heading.alignment = alignment
        # 减少标题后的间距
        heading.paragraph_format.space_after = Pt(6)
        return heading
    elif element_type == 'paragraph':
        p = doc.add_paragraph()
        run = p.add_run(text)
        run.font.name = font_name
        run._element.rPr.rFonts.set(qn('w:eastAsia'), font_name)
        run.font.size = Pt(font_size)
        run.font.bold = bold
        if color:
            run.font.color.rgb = RGBColor(*color)
        p.alignment = alignment
        # 减少段落间距
        p.paragraph_format.space_after = Pt(3)
        return p

def create_beautiful_table(doc, headers, data):
    """创建美观的表格，所有文字上下左右居中"""
    table = doc.add_table(rows=1, cols=len(headers))
    table.style = 'Table Grid'
    table.alignment = WD_TABLE_ALIGNMENT.CENTER
    
    # 设置表头样式
    hdr_cells = table.rows[0].cells
    for i, header in enumerate(headers):
        hdr_cells[i].text = header
        # 设置表头样式 - 所有表头居中
        for paragraph in hdr_cells[i].paragraphs:
            paragraph.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
            for run in paragraph.runs:
                run.font.bold = True
                run.font.size = Pt(10)
                run.font.color.rgb = RGBColor(255, 255, 255)  # 白色文字
        # 设置表头背景色和垂直居中
        hdr_cells[i]._tc.get_or_add_tcPr().append(parse_xml(f'<w:shd xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main" w:fill="4472C4"/>'))
        hdr_cells[i]._tc.get_or_add_tcPr().append(parse_xml(f'<w:vAlign xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main" w:val="center"/>'))
    
    # 添加数据行 - 所有数据上下左右居中
    for row_data in data:
        row_cells = table.add_row().cells
        for i, cell_data in enumerate(row_data):
            row_cells[i].text = str(cell_data)
            # 设置单元格样式 - 所有列都居中
            for paragraph in row_cells[i].paragraphs:
                paragraph.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
                for run in paragraph.runs:
                    run.font.size = Pt(9)
                    if i == 1:  # 分数列加粗
                        run.font.bold = True
                        run.font.color.rgb = RGBColor(68, 114, 196)  # 蓝色
            # 设置垂直居中
            row_cells[i]._tc.get_or_add_tcPr().append(parse_xml(f'<w:vAlign xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main" w:val="center"/>'))
    
    return table

def json2word(json_data, output_format='pdf'):
    """生成Word文档或PDF报告"""
    # 创建美观的图表
    create_beautiful_charts(json_data)
    
    # 创建 Word 文档
    doc = Document()
    
    # 设置页面边距 - 减小边距减少空行，让内容更紧凑
    sections = doc.sections
    for section in sections:
        section.top_margin = Inches(0.5)
        section.bottom_margin = Inches(0.5)
        section.left_margin = Inches(0.5)
        section.right_margin = Inches(0.5)
    
    # 设置默认字体
    doc.styles['Normal'].font.name = u'宋体'
    doc.styles['Normal']._element.rPr.rFonts.set(qn('w:eastAsia'), u'宋体')
    doc.styles['Normal'].font.size = Pt(12)
    
    # 添加文档标题
    add_element_with_style(doc, 'heading', "心理测评报告", 18, '黑体', True, 
                          (68, 114, 196), WD_PARAGRAPH_ALIGNMENT.CENTER)
    
    # 添加分隔线
    doc.add_paragraph("=" * 40).alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
    
    # 添加总分信息（带背景色）
    total_score_para = doc.add_paragraph()
    total_score_para.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
    run = total_score_para.add_run(f"总分：{json_data['all_score']} 分")
    run.font.name = u'微软雅黑'
    run._element.rPr.rFonts.set(qn('w:eastAsia'), u'微软雅黑')
    run.font.size = Pt(14)
    run.font.bold = True
    run.font.color.rgb = RGBColor(255, 255, 255)
    
    # 为总分段落添加背景色
    total_score_para._p.get_or_add_pPr().append(parse_xml(f'<w:shd xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main" w:fill="4472C4"/>'))
    
    # 插入图表部分 - 减少空行，让内容更紧凑
    add_element_with_style(doc, 'heading', "各题得分分布图", 14, '黑体', True, (68, 114, 196))
    # 插入图片并居中
    pic_para = doc.add_paragraph()
    pic_para.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
    pic_para.add_run().add_picture("all_score.png", width=Inches(6))
    # 减少图片后的间距
    pic_para.paragraph_format.space_after = Pt(3)
    
    add_element_with_style(doc, 'heading', "各维度得分图", 14, '黑体', True, (68, 114, 196))
    # 插入图片并居中
    pic_para2 = doc.add_paragraph()
    pic_para2.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
    pic_para2.add_run().add_picture("dimensions.png", width=Inches(6))
    # 减少图片后的间距
    pic_para2.paragraph_format.space_after = Pt(3)
    
    # 移除分页符，让内容连续显示
    # doc.add_page_break()  # 注释掉这行
    
    # 处理评估数据
    eval_data = json_data["eval"]
    eval_data_str = json.dumps(eval_data, ensure_ascii=False, indent=4)
    
    # 调用大模型优化内容
    prompt_t = f"""请优化以下JSON数据中的字段内容：
    结合name，text，help，score综合描述，仅优化列表中每个元素的text和help字段的文本表达，使其：
    表达更清晰准确；
    语言更加专业、规范；
    建议更具体、实用，适合家长、教师、心理工作者阅读；
    保留原有信息逻辑和内容层次，不新增、删除信息；
    禁止更改JSON结构和字段名称，禁止更改数值或其他字段内容。
    请仅返回优化后的完整JSON，保持格式不变。
    以下是需要优化的JSON数据：

    {eval_data_str}
    """
    print(prompt_t)
    print("大模型生成中")
    completion = client.chat.completions.create(
        model="qwen3-32b",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt_t},
        ],
        extra_body={"enable_thinking": False},
    )
    print(completion.choices[0].message.content)
    raw_content = completion.choices[0].message.content
    json_data = clean_llm_response(raw_content)
    print("清理后的数据：")
    print(json_data)

    # 创建评估结果表格
    add_element_with_style(doc, 'heading', "评估结果与建议", 16, '黑体', True, (68, 114, 196))
    
    # 准备表格数据
    table_headers = ['维度', '分数', '评估结果', '建议']
    table_data = []
    
    for item in json_data:
        if item["name"] == "总分":
            continue
        table_data.append([
            item["name"],
            str(item["score"]),
            item["text"].replace("\n", ""),
            item["help"]
        ])
    
    # 创建美观表格
    create_beautiful_table(doc, table_headers, table_data)
    
    # 生成总结
    summary_prompt = f"""
    请根据以下内容，生成一段简明、专业的总结性分析，适合直接写入心理测评报告中，语言要求：
    - 用正式、客观、规范表述；
    - 面向家长、教师、心理工作者；
    - 内容应简洁明了，避免重复冗余；
    - 字数在300字左右；
    - 总结需清晰指出儿童整体焦虑风险水平，并提供简要建议；
    - 只输出总结段落本身，不要输出其他说明或格式符号。

    以下为总分解释内容：
    {json_data}
    """
    print("大模型生成总结中...")
    completion_summary = client.chat.completions.create(
        model="qwen3-32b",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": summary_prompt},
        ],
        extra_body={"enable_thinking": False},
    )
    summary_content = completion_summary.choices[0].message.content.strip()
    print("生成的总结：")
    print(summary_content)
    
    # 添加总结部分
    add_element_with_style(doc, 'heading', "总结分析", 16, '黑体', True, (68, 114, 196))
    
    # 为总结添加特殊样式，并设置首行缩进为中文两个字符
    summary_para = doc.add_paragraph() 
    summary_para.alignment = WD_PARAGRAPH_ALIGNMENT.JUSTIFY
    summary_para.paragraph_format.first_line_indent = Pt(28)  # 中文2字符缩进，2*14pt=28pt
    summary_run = summary_para.add_run(summary_content)
    summary_run.font.name = u'微软雅黑'
    summary_run._element.rPr.rFonts.set(qn('w:eastAsia'), u'微软雅黑')
    summary_run.font.size = Pt(11)
    summary_run.font.color.rgb = RGBColor(51, 51, 51)
    
    # 为总结段落添加边框和背景
    summary_para._p.get_or_add_pPr().append(parse_xml(f'<w:pBdr xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:top w:val="single" w:sz="4" w:space="0" w:color="4472C4"/><w:left w:val="single" w:sz="4" w:space="0" w:color="4472C4"/><w:bottom w:val="single" w:sz="4" w:space="0" w:color="4472C4"/><w:right w:val="single" w:sz="4" w:space="0" w:color="4472C4"/></w:pBdr>'))
    summary_para._p.get_or_add_pPr().append(parse_xml(f'<w:shd xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main" w:fill="F8F9FA"/>'))
    
    # 添加页脚信息
    footer_para = doc.add_paragraph("报告生成时间：" + datetime.now().strftime("%Y年%m月%d日 %H:%M:%S"))
    footer_para.alignment = WD_PARAGRAPH_ALIGNMENT.RIGHT
    footer_run = footer_para.runs[0]
    footer_run.font.size = Pt(9)
    footer_run.font.color.rgb = RGBColor(128, 128, 128)
    
    # 根据输出格式选择保存方式
    if output_format.lower() == 'pdf':
        # 保存为Word文档然后转换为PDF
        doc.save("psychology_report.docx")
        try:
            convert("psychology_report.docx", "psychology_report.pdf")
            print("美观的连续PDF报告已生成！")
            return "psychology_report.pdf"
        except Exception as e:
            print(f"PDF转换失败: {e}")
            print("已生成Word文档作为备选")
            return "psychology_report.docx"
    else:
        # 保存Word文档
        doc.save("psychology_report.docx")
        print("美观的连续Word报告已生成！")
        return "psychology_report.docx"

def process_json(data):
    # 深拷贝避免破坏原数据
    new_eval = copy.deepcopy(data["eval"])
    # 总分，添加 score
    new_eval[0]["score"] = data["all_score"]
    # 遍历 eval 从第二条开始
    for i in range(1, len(new_eval)):
        # dimensions 的索引是 i-1
        dim = data["dimensions"][i - 1]
        # 替换 name
        new_eval[i]["name"] = dim["name"]
        # 添加 score
        new_eval[i]["score"] = dim["score"]
    # 更新回 json_data
    data["eval"] = new_eval
    return data

def clean_llm_response(response):
    # 数据清理
    # 1. 去掉两侧多余的单引号
    clean_data = response.strip("'")
    # 2. 去掉 ```json 和 ``` 部分
    if clean_data.startswith("```json"):
        clean_data = clean_data[len("```json"):]
    if clean_data.endswith("```"):
        clean_data = clean_data[:-3]
    # 3. 去掉开头和结尾的换行符
    clean_data = clean_data.strip()
    # 4. 反序列化为Python对象
    response = json.loads(clean_data)
    # 5. 打印结果
    # print(json_data)
    return response



app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 或者写局域网内允许的地址列表
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ✅ 定义请求体结构
class JsonPayload(BaseModel):
    data: Any  # 或具体定义字段：name: str, age: int 这样也可以
    output_format: str = 'pdf'  # 新增输出格式参数

# 新增输出格式参数，使用pdf格式可以得到更好的阅读体验
@app.post("/get_doc")
def get_doc(post_data: JsonPayload):
    process_data = process_json(post_data.data)
    output_file = json2word(process_data, post_data.output_format)
    
    # 根据输出文件类型返回相应的响应
    if output_file.endswith('.pdf'):
        return FileResponse(
            path="./psychology_report.pdf",
            media_type='application/pdf',
            filename='心理测评报告.pdf'
        )
    else:
        return FileResponse(
            path="./psychology_report.docx",
            media_type='application/vnd.openxmlformats-officedocument.wordprocessingml.document',
            filename='心理测评报告.docx'
        )

if __name__ == '__main__':
    # 假设 json 数据
    json_data = {
        "all_score": 24,
        "score_all": [
            0,
            2,
            0,
            0,
            0,
            1,
            1,
            1,
            0,
            0,
            1,
            1,
            0,
            1,
            0,
            0,
            1,
            1,
            1,
            0,
            1,
            1,
            1,
            0,
            0,
            0,
            1,
            0,
            1,
            0,
            1,
            1,
            0,
            1,
            1,
            2,
            0,
            1,
            1,
            0,
            0
        ],
        "dimensions": [
            {
                "name": "躯体化/惊恐因子",
                "score": 8
            },
            {
                "name": "广泛性焦虑因子",
                "score": 8
            },
            {
                "name": "分离性焦虑因子",
                "score": 4
            },
            {
                "name": "社交恐怖因子",
                "score": 2
            },
            {
                "name": "学校恐怖因子",
                "score": 2
            }
        ],
        "eval": [
            {
                "name": "总分",
                "text": "1.总分较低（＜25分）：如果儿童的SCARED量表总分低于25分，通常表明该儿童在当前阶段没有明显的焦虑性情绪障碍表现。但这并不意味着可以完全忽视儿童的情绪状态，家长和教师仍应关注儿童日常的情绪变化和行为表现，营造积极健康的成长环境，培养儿童良好的情绪调节能力。\n2.总分较高（≥25分）：当儿童总分达到或超过25分，说明该儿童可能存在焦虑性情绪障碍，需要引起重视。建议及时寻求专业心理评估机构或精神科医生的进一步诊断和评估，以便明确是否存在焦虑障碍以及确定具体类型，从而制定针对性的干预方案。在等待专业评估期间，家长和教师应密切关注孩子的情绪和行为变化，给予更多的关心和支持。",
                "help": ""
            },
            {
                "name": "维度1",
                "text": "躯体化/惊恐维度得分高：若该维度得分较高，说明儿童可能频繁出现与惊恐相关的躯体症状。",
                "help": "建议到医院进行全面的身体检查，排除生理疾病导致这些症状的可能性。同时，通过适当的运动、放松训练（如深呼吸、渐进性肌肉松弛等）来缓解身体的紧张感。例如，在躯体化/惊恐维度得分偏高，提示可能经常出现心跳加快、头晕等躯体不适。建议先去医院做身体检查，排查生理原因。日常可进行简单的放松练习，帮助其减轻身体的紧张状态。"
            },
            {
                "name": "维度2",
                "text": "广泛性焦虑维度得分高：对于此维度得分高的儿童，可能在日常生活中广泛存在的焦虑担忧。",
                "help": "可以通过与孩子一起制定计划、合理安排生活等方式，增强孩子对生活的掌控感，减少过度担忧。比如，在广泛性焦虑维度得分较高，表明其可能常常为多种事情担忧。家长和老师平时多与孩子沟通，引导他说出担忧的事情，帮助其制定合理的计划，逐步缓解焦虑情绪。"
            },
            {
                "name": "维度3",
                "text": "得分高表明儿童在与重要他人分离时存在较为严重的焦虑情绪。",
                "help": "如果儿童在这个维度得分高，家长要注意逐渐培养孩子的独立性，适当增加孩子与家人短暂分离的机会，让孩子逐渐适应。在分离时，给予孩子足够的安抚和信心。比如，分离性焦虑维度得分偏高，可能对与家人分离较为敏感。家长可逐步增加孩子独自活动的时间，分离前给予充分安慰，帮助孩子建立安全感，减少分离焦虑。"
            },
            {
                "name": "维度4",
                "text": "躯体化/惊恐维度得分高：若该维度得分较高，说明儿童可能频繁出现与惊恐相关的躯体症状。",
                "help": "社交恐怖维度得分高：针对社交恐怖维度得分高的儿童，建议家长和教师鼓励孩子多参与社交活动，从简单的小范围社交开始，如家庭聚会、小组活动等，逐步提升孩子的社交技能和自信心。例如，社交恐怖维度得分较高，提示其在社交场合可能存在紧张和回避行为。家长和老师可鼓励孩子从参与一些小型社交活动入手，帮助孩子积累成功的社交经验，提升社交自信。"
            },
            {
                "name": "维度5",
                "text": "躯体化/惊恐维度得分高：若该维度得分较高，说明儿童可能频繁出现与惊恐相关的躯体症状。",
                "help": "当儿童学校恐怖维度得分高时，家长和教师要与孩子深入沟通，了解其害怕学校的具体原因，如是否在学校遭遇了欺凌、学习困难等。针对具体问题，采取相应措施，如与学校老师沟通解决欺凌问题，为孩子提供学习辅导等。学校恐怖维度得分偏高，可能对学校存在恐惧心理。家长和老师需与孩子耐心沟通，找出害怕学校的根源，共同努力解决问题，帮助孩子重新建立对学校的积极态度。"
            }
        ]
    }
    # 调用方法如下：
    # data = process_json(json_data)
    # 默认生成PDF格式，避免空白问题
    # output_file = json2word(data, 'pdf')
    # print(f"报告已生成: {output_file}")

    # uvicorn.run(app, host="192.168.1.101", port=8001)
    uvicorn.run(app, host="127.0.0.1", port=8001)
    # uvicorn.run(app, host="172.20.0.89", port=8001)

